#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复ChromaDB路径配置问题
统一所有ChromaDB数据到配置文件指定的路径
"""

import os
import shutil
import yaml
from pathlib import Path

def load_config():
    """加载配置文件"""
    config_path = Path('config/database.yaml')
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def get_directory_size(path):
    """获取目录大小（MB）"""
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(path):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                if os.path.exists(filepath):
                    total_size += os.path.getsize(filepath)
    except Exception as e:
        print(f"计算目录大小失败 {path}: {e}")
    return total_size / (1024 * 1024)  # 转换为MB

def analyze_chromadb_directories():
    """分析现有的ChromaDB目录"""
    project_root = Path.cwd()
    
    # 可能的ChromaDB目录
    possible_dirs = [
        project_root / 'chroma_db',
        project_root / 'data' / 'chroma_db',
        project_root / 'data' / 'chromadb'
    ]
    
    existing_dirs = []
    for dir_path in possible_dirs:
        if dir_path.exists():
            size_mb = get_directory_size(dir_path)
            file_count = sum(1 for _ in dir_path.rglob('*') if _.is_file())
            existing_dirs.append({
                'path': dir_path,
                'size_mb': size_mb,
                'file_count': file_count
            })
    
    return existing_dirs

def fix_chromadb_path():
    """修复ChromaDB路径配置"""
    print("=== 修复ChromaDB路径配置 ===")
    
    # 1. 加载配置
    config = load_config()
    configured_path = config['databases']['chromadb']['persist_directory']
    project_root = Path.cwd()
    target_path = project_root / configured_path.lstrip('./')
    
    print(f"配置文件指定路径: {configured_path}")
    print(f"目标绝对路径: {target_path}")
    
    # 2. 分析现有目录
    existing_dirs = analyze_chromadb_directories()
    
    if not existing_dirs:
        print("❌ 未找到任何ChromaDB目录")
        return False
    
    print("\n现有ChromaDB目录:")
    for i, dir_info in enumerate(existing_dirs):
        print(f"  {i+1}. {dir_info['path']}")
        print(f"     大小: {dir_info['size_mb']:.2f}MB")
        print(f"     文件数: {dir_info['file_count']}")
    
    # 3. 找到最大的目录作为主目录
    main_dir = max(existing_dirs, key=lambda x: x['size_mb'])
    main_path = main_dir['path']
    
    print(f"\n选择主目录: {main_path} (大小: {main_dir['size_mb']:.2f}MB)")
    
    # 4. 如果主目录不是目标路径，则迁移数据
    if main_path != target_path:
        print(f"\n需要迁移数据从 {main_path} 到 {target_path}")
        
        # 确保目标目录的父目录存在
        target_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 如果目标路径已存在，先备份
        if target_path.exists():
            backup_path = target_path.with_suffix('.backup')
            print(f"备份现有目标目录到: {backup_path}")
            if backup_path.exists():
                shutil.rmtree(backup_path)
            shutil.move(str(target_path), str(backup_path))
        
        # 移动主目录到目标路径
        print(f"移动 {main_path} -> {target_path}")
        shutil.move(str(main_path), str(target_path))
        
        # 合并其他目录的数据
        for dir_info in existing_dirs:
            other_path = dir_info['path']
            if other_path != main_path and other_path.exists():
                print(f"合并数据从 {other_path}")
                
                # 遍历其他目录中的所有文件和子目录
                for item in other_path.rglob('*'):
                    if item.is_file():
                        # 计算相对路径
                        rel_path = item.relative_to(other_path)
                        target_file = target_path / rel_path
                        
                        # 确保目标目录存在
                        target_file.parent.mkdir(parents=True, exist_ok=True)
                        
                        # 如果目标文件不存在，则复制
                        if not target_file.exists():
                            shutil.copy2(str(item), str(target_file))
                            print(f"  复制: {rel_path}")
                
                # 删除已合并的目录
                print(f"删除已合并的目录: {other_path}")
                shutil.rmtree(other_path)
    
    else:
        print(f"\n✅ 主目录已经是目标路径，无需迁移")
        
        # 仍需要清理其他目录
        for dir_info in existing_dirs:
            other_path = dir_info['path']
            if other_path != target_path and other_path.exists():
                print(f"删除重复目录: {other_path}")
                shutil.rmtree(other_path)
    
    # 5. 验证最终结果
    if target_path.exists():
        final_size = get_directory_size(target_path)
        final_count = sum(1 for _ in target_path.rglob('*') if _.is_file())
        print(f"\n✅ 修复完成!")
        print(f"最终ChromaDB目录: {target_path}")
        print(f"大小: {final_size:.2f}MB")
        print(f"文件数: {final_count}")
        return True
    else:
        print(f"\n❌ 修复失败，目标目录不存在: {target_path}")
        return False

if __name__ == '__main__':
    try:
        success = fix_chromadb_path()
        if success:
            print("\n🎉 ChromaDB路径配置修复成功!")
            print("现在所有ChromaDB操作都将使用统一的数据仓库。")
        else:
            print("\n❌ ChromaDB路径配置修复失败!")
    except Exception as e:
        print(f"\n💥 修复过程中发生错误: {e}")
        import traceback
        traceback.print_exc()